import asyncio
import logging
import os
from typing import AsyncGenerator

import httpx
from mcp.client.streamable_http import streamable_http_client
from strands import Agent
from strands.tools.mcp import MCPClient
from strands.agent.agent_result import AgentResult
from strands_tools import current_time, http_request, calculator, python_repl
from query.tools.external_tools import search_pubmed, search_patents_web
from strands.agent.conversation_manager import SlidingWindowConversationManager
from strands.session.file_session_manager import FileSessionManager


from annotation.timer_annotation import log_time
from query.util.agent_debug import debugger_callback_handler
from query.model.tool_request_enum import ToolRequestEnum
from query.model.agent_metrics import AgentMetrics

logger = logging.getLogger(__name__)


# Hard caps — see Stage 1 Assessment §04.2 (Hanging Runs root causes)
_AGENT_MAX_ITERATIONS = int(os.environ.get("EUGENE_AGENT_MAX_ITERATIONS", "20"))
_AGENT_STREAM_TIMEOUT_S = float(os.environ.get("EUGENE_AGENT_STREAM_TIMEOUT_S", "180"))
_MCP_VERIFY_SSL = os.environ.get("EUGENE_MCP_VERIFY_SSL", "true").lower() != "false"
_ALLOW_PYTHON_REPL = os.environ.get("EUGENE_ALLOW_PYTHON_REPL", "false").lower() == "true"
# `calculator`/`current_time` create ReAct-loop surface area for questions that
# have nothing to do with arithmetic or clock. They are opt-in per request.
_ALLOW_UTILITY_TOOLS = os.environ.get("EUGENE_ALLOW_UTILITY_TOOLS", "false").lower() == "true"
# Per-tool-call duplicate-input cap — if the agent issues the same
# (tool_name, input_json) more than this many times, we raise.
_MAX_DUPLICATE_TOOL_CALLS = 2
# Absolute hard cap on any single stream's tool calls.
_HARD_TOOL_CALL_BUDGET = int(os.environ.get("EUGENE_MAX_TOOL_CALLS", "25"))


class EugeneDataAgent:
    system_prompt = """You are Eugene, an agent that answers biomedical competitive-intelligence questions by querying the CSL Behring Eugene knowledge graph.
Scope: drugs, diseases, gene/proteins, pathways, patents, clinical trials, organizations, and the relationships between them.

Follow the ReAct pattern (Reason → Act → Observe). Prefer MCP tools (fetch_*, find_*, lookup_*) for facts. When the graph does not contain the requested information, state that clearly and stop — do not loop.

TOOL GUIDE:
- Graph lookup (drugs, genes, diseases, organizations, patents already ingested): use the MCP `fetch_*` / `lookup_*` / `find_organization_*` tools.
- Latest patents on the open web: call `search_patents_web(query)` — it returns a Google Patents deep-link the user can follow. Do NOT try to scrape Google Patents with `http_request`, it will return 503.
- Biomedical literature (PubMed papers, including patent-adjacent publications): call `search_pubmed(query, max_results)`. It uses NCBI E-utilities and returns structured {pmid, title, journal, year, authors, url} records.
- Arbitrary URL fetch: `http_request` — only for non-biomedical pages. Many biomedical sites rate-limit bots.

HARD RULES:
1. Do NOT call `calculator` or `current_time` unless the user's question is explicitly about arithmetic or the current date/time. They are forbidden for data lookups.
2. Do NOT retry the same tool with the same arguments more than twice. If two attempts fail, report the failure and move on.
3. Do NOT invent data. If no tool returns an answer, say "I don't have this in the Eugene knowledge graph — here is a web search link instead" and include the search_patents_web or search_pubmed output.
4. Keep total tool calls to <= 20 per user turn. Prefer fewer, more precise tool calls.

Citation requirement: when you assert a fact from a tool result, cite the Eugene node_id(s) and tool used in the form [node_id=<id>, tool=<tool_name>] so the UI can render a context graph of your reasoning path.
""".strip()

    def __init__(self, eugene_mcp_server_url: str, model):
        self.eugene_mcp_server_url = eugene_mcp_server_url
        self.model = model

    def _new_conversation_manager(self) -> SlidingWindowConversationManager:
        # AGT-01 fix: instantiate per-conversation, not shared across concurrent sessions
        return SlidingWindowConversationManager(
            window_size=10, should_truncate_results=True, per_turn=2
        )

    @log_time
    def execute(
        self,
        token: str,
        user_prompt: str,
        conversation_id: str,
        include_tools: list[ToolRequestEnum] = [],
    ) -> AgentResult:
        mcp_client = self._build_mcp_client(token=token)
        agent = self._init_agent(
            mcp_client=mcp_client,
            conversation_id=conversation_id,
            include_tools=include_tools,
        )
        if logger.isEnabledFor(logging.INFO):
            logger.info(f"conversation: {conversation_id}")
            logger.info(f"model: {self.model}")
            logger.info(f"user prompt: {user_prompt}")
            logger.info(f"tools: {agent.tool_names}")

        with mcp_client:
            response = agent(user_prompt)

        if logger.isEnabledFor(logging.INFO):
            metrics = AgentMetrics.of(response)
            logger.info(f"{conversation_id} response metrics: {metrics}")
        return response

    @log_time
    async def execute_stream(
        self,
        token: str,
        user_prompt: str,
        conversation_id: str,
        include_tools: list[ToolRequestEnum] = [],
    ) -> AsyncGenerator[dict[str, str], None]:
        try:
            logger.info(f"Starting stream for conversation {conversation_id}")

            mcp_client = self._build_mcp_client(token=token)
            agent = self._init_agent(
                mcp_client=mcp_client,
                conversation_id=conversation_id,
                include_tools=include_tools,
            )

            yield {"type": "session", "session_id": conversation_id, "content": ""}

            # AGT-03 fix: wrap entire stream in a timeout so a stalled MCP/LLM can't hang the connection
            async def _run() -> AsyncGenerator[dict, None]:
                seen_tool_ids: set[str] = set()
                tool_call_budget = {"remaining": _HARD_TOOL_CALL_BUDGET}
                dup_fingerprints: dict[str, int] = {}

                def _register_tool_call(ev: dict) -> str | None:
                    """Return a breaker message if the loop should be killed."""
                    if ev.get("type") != "tool_call":
                        return None
                    tool_call_budget["remaining"] -= 1
                    if tool_call_budget["remaining"] < 0:
                        return (
                            f"Tool-call budget exhausted ({_HARD_TOOL_CALL_BUDGET}). "
                            "Stopping to prevent runaway ReAct loop."
                        )
                    import json as _json
                    try:
                        fp = f"{ev.get('tool')}::{_json.dumps(ev.get('tool_input'), sort_keys=True, default=str)}"
                    except Exception:
                        fp = f"{ev.get('tool')}::<unhashable>"
                    dup_fingerprints[fp] = dup_fingerprints.get(fp, 0) + 1
                    if dup_fingerprints[fp] > _MAX_DUPLICATE_TOOL_CALLS:
                        return (
                            f"Agent repeated `{ev.get('tool')}` with identical inputs "
                            f"{dup_fingerprints[fp]} times — breaking loop."
                        )
                    return None

                with mcp_client:
                    kill_reason: str | None = None
                    async for event in agent.stream_async(user_prompt):
                        # Text deltas
                        if isinstance(event, dict) and "data" in event:
                            yield {
                                "type": "content",
                                "content": event["data"],
                                "session_id": conversation_id,
                            }
                        # Opportunistic mid-stream extraction (best-effort, shape may vary)
                        for ev in _extract_tool_events(event, conversation_id, seen_tool_ids):
                            yield ev
                            kill_reason = _register_tool_call(ev)
                            if kill_reason:
                                break
                        if kill_reason:
                            break

                    if kill_reason:
                        logger.warning(f"{conversation_id}: {kill_reason}")
                        yield {
                            "type": "content",
                            "content": f"\n\n⚠️ {kill_reason}",
                            "session_id": conversation_id,
                        }

                    # Post-stream: authoritative tool-call extraction from the
                    # agent's full message history. This is the reliable path
                    # because Strands' in-flight event shapes differ across
                    # versions/providers, but `agent.messages` is stable.
                    for ev in _harvest_messages(agent, conversation_id, seen_tool_ids):
                        yield ev

            async def _collect():
                async for item in _run():
                    yield item

            try:
                async for item in self._with_timeout(_collect(), _AGENT_STREAM_TIMEOUT_S):
                    yield item
            except asyncio.TimeoutError:
                logger.warning(
                    f"conversation {conversation_id} exceeded {_AGENT_STREAM_TIMEOUT_S}s timeout"
                )
                yield {
                    "type": "error",
                    "content": f"Agent execution exceeded {int(_AGENT_STREAM_TIMEOUT_S)}s timeout",
                    "session_id": conversation_id,
                }

            yield {"type": "done", "content": "", "session_id": conversation_id}
            logger.info(f"conversation {conversation_id}: completed streaming")

        except Exception as e:
            logger.exception(
                f"Error while streaming conversation {conversation_id}: {e}"
            )
            yield {
                "type": "error",
                "content": f"Error: {str(e)}",
                "session_id": conversation_id,
            }

    @staticmethod
    async def _with_timeout(agen: AsyncGenerator, timeout_s: float):
        """Yield from `agen` but bound total wall time to `timeout_s`."""
        loop = asyncio.get_running_loop()
        deadline = loop.time() + timeout_s
        while True:
            remaining = deadline - loop.time()
            if remaining <= 0:
                raise asyncio.TimeoutError()
            try:
                item = await asyncio.wait_for(agen.__anext__(), timeout=remaining)
            except StopAsyncIteration:
                return
            yield item

    def _init_agent(
        self,
        mcp_client: MCPClient,
        conversation_id: str,
        include_tools: list[ToolRequestEnum] = [],
    ) -> Agent:
        logger.info(f"include_tools={include_tools}")
        name = "EugeneDataAgent"

        # Must enter the MCP client context to list tools; the caller re-enters it
        # for the actual agent run, which is safe because MCPClient is re-entrant
        # via its internal session management.
        with mcp_client:
            session_manager = FileSessionManager(session_id=conversation_id)
            logger.info(f"writing session to {session_manager.storage_dir}")

            # Utility tools are opt-in — they caused runaway ReAct loops
            # (agent spamming `current_time` trying to "search the web").
            # See Stage 1 Assessment §04.2 (Hanging Runs). AGT-05 for python_repl.
            local_tools: list = []
            if _ALLOW_UTILITY_TOOLS:
                local_tools.extend([calculator, current_time])
            if _ALLOW_PYTHON_REPL:
                local_tools.append(python_repl)
            if ToolRequestEnum.HTTP in include_tools:
                # Real web-backed tools (PubMed E-utilities + Google Patents link builder).
                # These are preferred over raw http_request because they hit
                # bot-friendly APIs and return structured JSON the UI can render.
                local_tools.extend([search_pubmed, search_patents_web, http_request])

            mcp_tools = []
            if ToolRequestEnum.EUGENE in include_tools:
                mcp_tools.extend(mcp_client.list_tools_sync())
            tools = [*mcp_tools, *local_tools]

        agent = Agent(
            name=name,
            tools=tools,
            model=self.model,
            system_prompt=self.system_prompt,
            callback_handler=debugger_callback_handler,
            # AGT-01 fix: per-conversation manager
            conversation_manager=self._new_conversation_manager(),
            session_manager=session_manager,
        )
        # AGT-02 fix: cap ReAct iterations. Strands 1.21 does not expose a
        # constructor kwarg for this; set it post-construction defensively so
        # newer releases pick it up and older releases simply ignore it.
        for attr in ("max_iterations", "max_parallel_tools"):
            try:
                setattr(agent, attr, _AGENT_MAX_ITERATIONS if attr == "max_iterations" else 1)
            except Exception:
                pass
        logger.info(f"initialized {name}, tools={agent.tool_names}")
        return agent

    def _build_mcp_client(self, token: str) -> MCPClient:
        # SEC-01 fix: default verify=True; dev can opt out via EUGENE_MCP_VERIFY_SSL=false
        verify = _MCP_VERIFY_SSL
        return MCPClient(
            lambda: streamable_http_client(
                url=self.eugene_mcp_server_url,
                http_client=httpx.AsyncClient(
                    verify=verify,
                    headers={"Authorization": f"Bearer {token}"},
                    timeout=httpx.Timeout(connect=10.0, read=60.0, write=30.0, pool=10.0),
                ),
            )
        )


# ----------------------------------------------------------------------------
# Tool-call event extraction helpers
# ----------------------------------------------------------------------------

def _extract_tool_events(event, conversation_id: str, seen: set):
    """Yield tool_call / tool_result events if the mid-stream `event` contains them.

    Strands event shapes vary across providers (OpenAI vs Anthropic) and minor
    versions, so this is defensive: we inspect several known keys and skip on
    anything unexpected.
    """
    if not isinstance(event, dict):
        return

    # Path 1: `current_tool_use` (Anthropic-shaped streaming)
    cur = event.get("current_tool_use")
    if isinstance(cur, dict):
        tid = cur.get("toolUseId") or cur.get("id") or ""
        if tid and tid not in seen:
            seen.add(tid)
            yield {
                "type": "tool_call",
                "content": "",
                "session_id": conversation_id,
                "tool": cur.get("name", ""),
                "tool_input": cur.get("input", {}),
                "tool_id": tid,
            }

    # Path 2: completed `message` events containing toolUse / toolResult blocks
    msg = event.get("message")
    if isinstance(msg, dict):
        for block in (msg.get("content") or []):
            if not isinstance(block, dict):
                continue
            if "toolUse" in block:
                t = block["toolUse"] or {}
                tid = t.get("toolUseId") or t.get("id") or ""
                if tid and tid not in seen:
                    seen.add(tid)
                    yield {
                        "type": "tool_call",
                        "content": "",
                        "session_id": conversation_id,
                        "tool": t.get("name", ""),
                        "tool_input": t.get("input", {}),
                        "tool_id": tid,
                    }
            if "toolResult" in block:
                r = block["toolResult"] or {}
                tid = r.get("toolUseId") or ""
                yield {
                    "type": "tool_result",
                    "content": "",
                    "session_id": conversation_id,
                    "tool_id": tid,
                    "tool_output": _flatten_tool_output(r.get("content", [])),
                }


def _harvest_messages(agent, conversation_id: str, seen: set):
    """After stream completes, walk `agent.messages` and emit any tool calls we missed.

    This is the authoritative path: `agent.messages` is Strands' conversation
    history and always contains the full toolUse/toolResult trail.
    """
    messages = getattr(agent, "messages", None) or []
    pending_calls: dict[str, dict] = {}
    for m in messages:
        if not isinstance(m, dict):
            continue
        for block in m.get("content") or []:
            if not isinstance(block, dict):
                continue
            if "toolUse" in block:
                t = block["toolUse"] or {}
                tid = t.get("toolUseId") or t.get("id") or ""
                if not tid:
                    continue
                pending_calls[tid] = {
                    "type": "tool_call",
                    "content": "",
                    "session_id": conversation_id,
                    "tool": t.get("name", ""),
                    "tool_input": t.get("input", {}),
                    "tool_id": tid,
                }
            if "toolResult" in block:
                r = block["toolResult"] or {}
                tid = r.get("toolUseId") or ""
                # Emit the paired tool_call first if we haven't already
                if tid and tid not in seen and tid in pending_calls:
                    seen.add(tid)
                    yield pending_calls[tid]
                yield {
                    "type": "tool_result",
                    "content": "",
                    "session_id": conversation_id,
                    "tool_id": tid,
                    "tool_output": _flatten_tool_output(r.get("content", [])),
                }
    # Any tool_call we discovered that never had a result (rare) — still emit
    for tid, call in pending_calls.items():
        if tid not in seen:
            seen.add(tid)
            yield call


def _flatten_tool_output(content):
    """Tool result `content` is a list of blocks like [{"text": "..."}, {"json": {...}}].
    Parse JSON strings so the UI can walk them for nodes/rels.
    """
    if not isinstance(content, list):
        return content
    out = []
    for block in content:
        if not isinstance(block, dict):
            out.append(block)
            continue
        if "json" in block:
            out.append(block["json"])
        elif "text" in block:
            text = block["text"]
            # Try to parse as JSON; fall back to raw text
            try:
                import json
                out.append(json.loads(text))
            except Exception:
                out.append(text)
        else:
            out.append(block)
    return out

